use std::collections::HashMap;
use std::cell::Cell;

type Monkey = &'static str;

#[derive(Debug)]
struct Expr {
    value: Cell<Option<i64>>,
    expr_kind: ExprKind,
}

impl Expr {
    fn new(expr_kind: ExprKind) -> Self {
        Self {
            value: Cell::new(None),
            expr_kind
        }
    }
}

#[derive(Debug)]
enum ExprKind {
    Number(i64),
    BinaryExpr(BinaryExpr)
}

#[derive(Debug)]
enum BinaryExpr {
    Add(Monkey, Monkey),
    Sub(Monkey, Monkey),
    Mul(Monkey, Monkey),
    Div(Monkey, Monkey)
}

fn parse_input() -> HashMap<Monkey, Expr> {
    include_str!("../input.txt")
        .lines()
        .map(|line| -> (Monkey, Expr) {
            let name = &line[..4];
            if line.chars().nth(6).unwrap().is_digit(10) {
                (name, Expr::new(ExprKind::Number(line[6..].parse().unwrap())))
            } else {
                let left_name = &line[6..10];
                let right_name = &line[13..17];
                match &line[11..12] {
                    "+" => (name, Expr::new(ExprKind::BinaryExpr(BinaryExpr::Add(left_name, right_name)))),
                    "-" => (name, Expr::new(ExprKind::BinaryExpr(BinaryExpr::Sub(left_name, right_name)))),
                    "*" => (name, Expr::new(ExprKind::BinaryExpr(BinaryExpr::Mul(left_name, right_name)))),
                    "/" => (name, Expr::new(ExprKind::BinaryExpr(BinaryExpr::Div(left_name, right_name)))),
                    _ => panic!("Bad input"),
                }
            }
        })
        .collect()
}

fn part1() {
    fn evaluate(expr: &Expr, exprs: &HashMap<Monkey, Expr>) -> i64 {
        let value = expr.value.take();
        match value {
            Some(v) => { v }
            None => {
                let value = calculate(expr, exprs);
                expr.value.set(Some(value));
                value
            }
        }
    }

    fn calculate(expr: &Expr, exprs: &HashMap<Monkey, Expr>) -> i64 {
        match &expr.expr_kind {
            ExprKind::Number(v) => *v,
            ExprKind::BinaryExpr(bin_expr) => {
                match bin_expr {
                    BinaryExpr::Add(l, r) => evaluate(&exprs[l], exprs) + evaluate(&exprs[r], exprs),
                    BinaryExpr::Sub(l, r) => evaluate(&exprs[l], exprs) - evaluate(&exprs[r], exprs),
                    BinaryExpr::Mul(l, r) => evaluate(&exprs[l], exprs) * evaluate(&exprs[r], exprs),
                    BinaryExpr::Div(l, r) => evaluate(&exprs[l], exprs) / evaluate(&exprs[r], exprs),
                }
            }
        }
    }

    let exprs = parse_input();
    let root = &exprs["root"];
    println!("Part 1: {}", evaluate(root, &exprs));
}

fn part2() {
    fn add(l: i64, r: i64) -> i64 { l + r }
    fn sub(l: i64, r: i64) -> i64 { l - r }
    fn mul(l: i64, r: i64) -> i64 { l * r }
    fn div(l: i64, r: i64) -> i64 { l / r }

    enum Calculator {
        Param,
        Number(i64),
        BinaryOper(fn(i64, i64) -> i64, Box<Calculator>, Box<Calculator>)
        
    }

    impl Calculator {
        fn evaluate(&self, input: i64) -> i64 {
            match self {
                Self::Param => input,
                Self::Number(v) => *v,
                Self::BinaryOper(func, left, right) => {
                    let lv = left.evaluate(input);
                    let rv = right.evaluate(input);
                    func(lv, rv)
                }
            }
        }
    }

    fn build_calculator(name: Monkey, exprs: &HashMap<Monkey, Expr>) -> Box<Calculator> {
        if name == "humn" {
            Box::new(Calculator::Param)
        } else {
            let expr = &exprs[name];
            match &expr.expr_kind {
                ExprKind::Number(v) => Box::new(Calculator::Number(*v)),
                ExprKind::BinaryExpr(bin_expr) => {
                    match bin_expr {
                        BinaryExpr::Add(l, r) => Box::new(Calculator::BinaryOper(add, build_calculator(l, exprs), build_calculator(r, exprs))),
                        BinaryExpr::Sub(l, r) => Box::new(Calculator::BinaryOper(sub, build_calculator(l, exprs), build_calculator(r, exprs))),
                        BinaryExpr::Mul(l, r) => Box::new(Calculator::BinaryOper(mul, build_calculator(l, exprs), build_calculator(r, exprs))),
                        BinaryExpr::Div(l, r) => Box::new(Calculator::BinaryOper(div, build_calculator(l, exprs), build_calculator(r, exprs))),
                    }
                }
            }
        }
    }

    let exprs = parse_input();
    let root = &exprs["root"];
    let (root_left, root_right) = match &root.expr_kind {
        ExprKind::Number(_) => panic!(""),
        ExprKind::BinaryExpr(bin_expr) => {
            match bin_expr {
                BinaryExpr::Add(l, r) => (build_calculator(l, &exprs), build_calculator(r, &exprs)),
                BinaryExpr::Sub(l, r) => (build_calculator(l, &exprs), build_calculator(r, &exprs)),
                BinaryExpr::Mul(l, r) => (build_calculator(l, &exprs), build_calculator(r, &exprs)),
                BinaryExpr::Div(l, r) => (build_calculator(l, &exprs), build_calculator(r, &exprs)),
            }
        }
    };

    for input in 0.. {
        if root_left.evaluate(input) == root_right.evaluate(input) {
            println!("Part 2: {}", input);
            break;
        }
    }

}

fn main() {
    part1();
    part2();
}
